﻿using System;
using System.Collections.Generic;
using System.Linq;
using Algoritmia.DataStructures.Digraphs;
using Algoritmia.DataStructures.Maps;
using Algoritmia.Utils;

namespace Algoritmia.Problems.ShortestPaths.KEdges
{

    public class KEdgesDistanceFinder<T> : IMinimumDistanceFinder<T,double>
    {
        private readonly Func<IEnumerable<T>, int, IMap<Tuple<T, int>, double>> _createMap;
        public int K { get; set; }

        public KEdgesDistanceFinder(int k, Func<IEnumerable<T>, int, IMap<Tuple<T, int>, double>> createMap = null)
        {
            K = k;
            _createMap = createMap ?? ((it, a) => AlgoritmiaFactory.Maps<Tuple<T, int>, double>.Default);
        }

        #region Implementation of IMinimumDistanceFinder<T,double>

        public double SomeToSomeDistance(IDigraph<T> g, WeightingFunction<T, double> d, IEnumerable<T> sources, IEnumerable<T> targets)
        {
            var distances = _createMap(g.V, K);

            foreach (var v in g.V)
            {
                distances[new Tuple<T, int>(v, 0)] = double.PositiveInfinity;
            }
            foreach (var s in sources)
            {
                distances[new Tuple<T, int>(s, 0)] = 0;
            }
            for (int i = 1; i <= K; i++)
            {
                foreach (var v in g.V)
                {
                    T u = v;
                    int c = i - 1;
                    distances[new Tuple<T, int>(v, i)] = (from p in g.Predecessors(v)
                                                          select distances[new Tuple<T, int>(p, c)] + d(p, u)).DefaultIfEmpty(double.PositiveInfinity).Min();
                }
            }
            return (from t in targets
                    select distances[new Tuple<T, int>(t, K)]).Min();
        }

        #endregion
    }

    public class SpaceReducedKEdgesDistanceFinder<T> : IMinimumDistanceFinder<T,double>
    {
        private readonly Func<IEnumerable<T>, IMap<T, double>> _createMap;
        public int K { get; set; }

        public SpaceReducedKEdgesDistanceFinder(int k, Func<IEnumerable<T>, IMap<T, double>> createMap = null)
        {
            K = k;
            _createMap = createMap ?? (it => AlgoritmiaFactory.Maps<T, double>.Default);
        }

        #region Implementation of IMinimumDistanceFinder<T,double>

        public double SomeToSomeDistance(IDigraph<T> g, WeightingFunction<T, double> d, IEnumerable<T> sources, IEnumerable<T> targets)
        {
            var prev = _createMap(g.V);
            var current = _createMap(g.V);
            foreach (var v in g.V)
            {
                current[v] = double.PositiveInfinity;
            }
            foreach (var s in sources)
            {
                current[s] = 0.0;
            }
            for (int i = 1; i <= K; i++)
            {
                var s = current;
                current = prev;
                prev = s;
                foreach (var v in g.V)
                {
                    var u = v;
                    current[v] = (from p in g.Predecessors(v)
                                  select prev[p] + d(p, u)).DefaultIfEmpty(double.PositiveInfinity).Min();
                }
            }
            return (from t in targets
                    select current[t]).Min();
        }

        #endregion
    }
}
